NVFP4 / AMD-FP4 outer-block (two-level) scaling for FP4 training

Background

NVFP4 quantizes tensors with a single per-tensor FP32 scale (tensorwise / PTS) on top of per-16-element FP4 (E2M1) blocks whose block scales live on the E4M3 grid. On LLM activations and MoE expert weights the magnitude distribution is heavy-tailed: a small fraction of channels/tokens ("massive activations") carry values orders of magnitude larger than the rest. With one global scale, those outliers force the per-block E4M3 inner scales of the ~99% non-outlier blocks toward underflow, collapsing their effective precision and degrading both forward-output and gradient SNR.

Outer-block (two-level) scaling inserts an intermediate per-super-block FP32 scale between the global level and the per-16 inner blocks. Each outer block (e.g. 64×64 or 1×128 tiles) is normalized by its own amax / (inner_grid_max · F4_E2M1_MAX) before inner-block FP4 quantization. This localizes each outlier to its own tile so the rest of the tensor keeps precision, directly targeting the PTS failure mode. The same machinery is reused for AMD-FP4, whose only difference is a wider UE5M3 inner-scale grid (max 114688 vs E4M3's 448).

Technical design

Three scale levels, applied per quantized operand (activations, weights, gradients) on both the forward-reduction and backward-reduction axes:

1. Outer-block FP32 scale — per super-block (1D 1×outer_block_size, or 2D outer_block_size×outer_block_size), computed as amax / (scale_max · 6), where scale_max is the inner grid max (E4M3 → 448 for NVFP4, UE5M3 → 114688 for AMD-FP4). Held in FP32 (no grid rounding).
2. Inner-block scale — the existing per-16 NVFP4/AMD-FP4 block scale, on the format's grid.
3. FP4 (E2M1) values — packed 2-per-byte as today.

Key properties:

• Tensorwise and outer-block are mutually exclusive (two_level_scaling ∈ {none, tensorwise, outer_block}).
• outer_block_size must be a multiple of the inner block (16) and divide every quantized reduction dim (GPT-OSS hidden=2880 is not divisible by 128, so the recipe uses 64).
• Independent 1D/2D layout knobs for activations and weights (use_outer_2dblock_x, use_outer_2dblock_w).
• For AMD-FP4 the format difference is absorbed by the (power-of-two) outer-scale denominator, so AMD-FP4 outer-block is numerically identical to NVFP4 outer-block run with scale_format="ue5m3" on well-ranged data; the two only diverge at intra-outer-block dynamic ranges wide enough to underflow one grid but not the other.

This PR delivers the core outer-block scaling for dense Linear and MoE grouped GEMM, in both NVFP4 and AMD-FP4. Two paper-recipe refinements (X-hat alignment and N-axis dX-RHT) were prototyped but are deferred / removed from this PR as not yet needed; they are preserved on a backup branch.

Implementation

• Quantization core (nvfp_quantization.py): convert_to/from_nvfp4_outer_block compute the per-outer-block FP32 scale, normalize, then call the existing inner NVFP4 kernel with tensorwise scaling disabled. The outer-scale denominator and the inner kernel both take scale_format, so one path serves E4M3 and UE5M3. _qdq gained use_outer_block_scale, is_2d_outer, outer_block_size, plus precomputed_outer_scale / return_outer_scale so an axis-invariant 2D outer scale is computed once and reused across the fprop (axis=-1) and wgrad (axis=0) QDQ calls.

• Dense linear (nvfp_linear.py, NVFP4LinearFunction): threads the outer-block knobs and the outer-scale-sharing cache through forward/backward; AMD-FP4 reuses it via _to_amdfp4_then_scaled_mm pinned to UE5M3.

• Grouped MoE GEMM (nvfp_grouped_gemm/*): per-expert/per-tile outer scale threaded through fprop, dgrad, and wgrad; AMD-FP4 grouped wrapper pins UE5M3.

• Dispatch (dispatch/tensor.py, config.py): routes outer-block kwargs to the dense and grouped helpers for both nvfp4 and amdfp4; TrainingOpConfig carries the outer-block fields.

• Recipe / validation (modifiers/lpt/base.py): LowPrecisionTrainingModifier accepts two_level_scaling="outer_block" for nvfp4 and amdfp4, validates outer_block_size, and rejects outer-side knobs outside outer-block mode.

• Recipes: alto/models/gpt_oss/configs/lpt_recipe_outer_block.yaml (NVFP4) and lpt_recipe_amdfp4_outer_block.yaml (AMD-FP4), e.g.:

  training_stage:
    lpt_modifiers:
      LowPrecisionTrainingModifier:
        scheme: "nvfp4"            # or "amdfp4"
        targets: ["Linear", "GptOssGroupedExperts"]
        ignore: ["output", "re:.*\.router\.gate"]
        two_level_scaling: outer_block
        use_outer_2dblock_w: true
        outer_block_size: 64
  

Op-level test coverage

New / extended suites (all GPU op-level, MI300X):

• Outer-block quantization (test_nvfp_quantization.py): outer-scale shape derivation for the four 1D/2D × inner/outer layouts; finite QDQ round-trips; bit-exact equality of the Triton path vs a PyTorch reference for convert_to/from_nvfp4_outer_block; PTS-vs-outer-block dominance on heterogeneous / sparse-outlier data; precomputed_outer_scale ↔ recompute equivalence; return_outer_scale plumbing; invalid-shape and PTS/outer-block mutual-exclusion guards; NaN-input sanitization.
• Dense linear outer-block (test_nvfp_outer_block_linear.py): autograd O/dX/dW SNR contract across the 1D/2D inner × 1D/2D outer matrix; independent X/W layouts; K=2048 SR stress; DGE-rejection guard; lognormal "massive activation" dominance over PTS (≥1 dB on O and dX); outer-block × M-axis Hadamard.
• Grouped MoE outer-block (test_nvfp_grouped_gemm_outer_block.py): per-expert/per-tile outer-block grouped GEMM op-level suite.
• AMD-FP4 outer-block (test_amdfp_linear.py): bit-identical parity of the AMD-FP4 dense outer-block wrapper to _to_nvfp4_then_scaled_mm(scale_format="ue5m3").
• Dispatch routing (test_nvfp_dispatch_guards.py, test_amdfp_dispatch_guards.py): NVFP4/AMD-FP4 linear and grouped_mm route to the right helper and propagate the outer-block kwargs.
• Recipe validation (test_outer_block_validation.py): outer-block accepted for nvfp4/amdfp4 and MoE targets; outer_block_size multiple-of-16 rule; non-outer-block misuse rejected.

Test status

• 901 / 901 op-level tests pass (tests/unittest/nvfp4, tests/unittest/amdfp4, alto/modifiers/lpt/tests) on MI300X (gfx942), torch 2.12.
• Regression check (no numerical drift): running the identical NVFP4 outer-block recipe on the GPT-OSS debug model with a fixed seed, pre-change vs post-change code produces bit-identical per-step train loss and grad-norm (Δ = 0.00000 at every step), confirming the integration introduces no regression on the outer-block path.

E2E results — GPT-OSS-20B + NVFP4 outer scale

All-layers FP4 (NVFP4 on MoE experts + attention Linears), 8× MI300X, MLPerf C4 data, 16,128 steps, identical optimizer / cosine-LR schedule. Reported loss is periodic validation loss (init → final):

Observations:

• No regression vs tensorwise: NVFP4 outer-block final loss (3.3438) matches PTS (3.3484) to within seed-to-seed noise (≤ 0.005) and is marginally lower, while exercising the full per-outer-block scale path on every MoE expert and Linear.
• Stable convergence: smooth 16k-step descent from ~12.65 with no divergence / NaN.
• Within ~1.6% of BF16: both FP4 recipes land ~0.055–0.059 above the BF16 reference (3.2897), the expected FP4 emulation gap; outer-block closes it slightly relative to PTS.

Scope / deferred

• Included: NVFP4 + AMD-FP4 outer-block scaling for dense Linear and MoE grouped GEMM, recipes, op-level tests, dispatch + recipe validation.
• Deferred (removed from this PR, kept on a backup branch): the TetraJet-style X-hat alignment and N-axis dX-RHT paper-recipe refinements. DGE + outer-block is explicitly guarded as unsupported.